import numpy as np
from paddle.nn.functional import binary_cross_entropy_with_logits, sigmoid, binary_cross_entropy
import paddle.nn as nn
import paddle


class loss(nn.Layer):
    def __init__(self, anchors, printloss=False, ignore_thresh=0.5):
        super(loss, self).__init__()
        self.anchors = anchors
        self.printloss = printloss
        self.ignore_thresh =ignore_thresh

    # feats中tx,ty,tw,th值的范围都为 0-1
    def head(self, feats, anchors, cal_loss=False):
        feats = feats.transpose((0, 2, 3, 1))
        anchors = paddle.to_tensor(anchors)
        batchsize = feats.shape[0]
        anchors = anchors.reshape((1, 1, 1, 3, 2))
        grid_shape = feats.shape[1:2][0]
        grid_y = paddle.tile(paddle.reshape(paddle.arange(0, grid_shape), (-1, 1, 1, 1)), (1, grid_shape, 1, 1))
        grid_x = paddle.tile(paddle.reshape(paddle.arange(0, grid_shape), [1, -1, 1, 1]), [grid_shape, 1, 1, 1])
        grid = paddle.concat([grid_x, grid_y], axis=-1)  # (13,13,1,2)
        feats = feats.reshape((batchsize, grid_shape, grid_shape, 3, 5+4))  # (bsize,13,13,3,9)
        box_xy = sigmoid(feats[..., :2]+grid) / grid_shape
        box_wh = paddle.exp(feats[..., 2:4]) * anchors / grid_shape
        box_confidence = sigmoid(feats[..., 4:5])
        box_class_probs = sigmoid(feats[..., 5:])
        if cal_loss is True:
            return grid, feats, box_xy, box_wh
        else:
            return box_xy, box_wh, box_confidence, box_class_probs

    def box_iou(self, b1, b2):
        '''
        b1: (...,4) xywh
        b2: (j,4)  xywh
        Returns
        iou  (...,j)
        '''
        b1 = paddle.unsqueeze(b1, -2)  # (...,1,4)
        b1_xy = b1[..., :2]
        b1_wh = b1[..., 2:4]
        b1_wh_half = b1_wh/2
        b1_mins = b1_xy - b1_wh_half
        b1_maxs = b1_xy + b1_wh_half

        b2 = paddle.unsqueeze(b2, 0)  # (1,j,4)
        b2_xy = b2[..., :2]
        b2_wh = b2[..., 2:4]
        b2_wh_half = b2_wh/2
        b2_mins = b2_xy - b2_wh_half
        b2_maxs = b2_xy + b2_wh_half
        inter_mins = paddle.maximum(b1_mins, b2_mins)  # (...,j,2)
        inter_maxs = paddle.minimum(b1_maxs, b2_maxs)
        temp = paddle.zeros_like(inter_maxs)
        inter_wh = paddle.maximum(inter_maxs-inter_mins, temp)
        inter_area = inter_wh[..., 0] * inter_wh[..., 1]  # (...,j)
        b1_area = b1_wh[..., 0] * b1_wh[..., 1]
        b2_area = b2_wh[..., 0] * b2_wh[..., 1]
        iou = inter_area/(b1_area+b2_area-inter_area)
        return iou

    def BinaryCrossEntropy(self, label, logit):
        out = paddle.max(logit, 0)-logit*label+paddle.log(1+paddle.exp(-paddle.abs(logit)))
        return out


    def forward(self, output, y_true):
        '''
        Args:
            output: (bs,13,13,3,5+num)
        Returns:
        '''
        anchors = np.array(self.anchors)
        batchsize = (output[0].shape)[0]
        anchors_mask = [[6,7,8], [3,4,5], [0,1,2]]
        input_shape = (416,416)
        grid_shape = [(13,13),(26,26),(52,52)]
        loss = 0
        for i in range(3):
            object_mask = y_true[i][..., 4:5]
            class_probs = y_true[i][..., 5:]

            grid, pred, pred_xy, pred_wh = self.head(feats=output[i], anchors=anchors[anchors_mask[i]], cal_loss=True)
            pred_box = paddle.concat([pred_xy, pred_wh], axis=-1)
            true_xy = y_true[i][..., :2]*paddle.to_tensor(grid_shape[i]) - grid
            # 处理不同
            true_wh = paddle.zeros_like(y_true[i][..., 2:4])
            true_wh[true_wh == 0] = 1
            true_wh = paddle.log(true_wh)/paddle.to_tensor(anchors[anchors_mask[i]]) * 416

            loss_scale = 2 - y_true[i][..., 2:3]*y_true[i][..., 3:4]

            ignore_mask_bool = object_mask.astype(bool)  # (bs,13,13,3,1)
            temp = paddle.tile(ignore_mask_bool, (1, 1, 1, 1, 4))  # (bs,13,13,3,4)
            true_box = y_true[i][..., 0:4]  # (bs,13,13,3,4)

            if len(true_box[temp]) == 0:
                true_box = paddle.to_tensor([0, 0, 0, 0], dtype=paddle.float32)
            else:
                true_box = true_box[temp].reshape((-1, 4))  # (j,4)
            # pre_box (bs,13,13,3,4)
            iou = self.box_iou(pred_box, true_box)  # (bs,13,13,3,j)
            best_iou = paddle.max(iou, axis=-1)  # (bs,13,13,3,1)
            # 为true则说明需要被忽略
            ignore_mask = best_iou < self.ignore_thresh

            xy_loss = object_mask*loss_scale*self.BinaryCrossEntropy(true_xy, pred[..., 0:2])
            wh_loss = object_mask*loss_scale*0.5*paddle.square(true_wh-pred[..., 2:4])
            ignore_mask = paddle.unsqueeze(ignore_mask, -1)
            confidence_loss = object_mask*self.BinaryCrossEntropy(object_mask, pred[..., 4:5])+(1-object_mask)*self.BinaryCrossEntropy(object_mask,pred[...,4:5])*ignore_mask
            class_loss = object_mask*self.BinaryCrossEntropy(class_probs,pred[..., 5:])
            xy_loss = xy_loss.sum()/batchsize
            wh_loss = wh_loss.sum()/batchsize
            confidence_loss = confidence_loss.sum()/batchsize
            class_loss = class_loss.sum()/batchsize
            loss += xy_loss+wh_loss+confidence_loss+class_loss
        if self.printloss:
            print('loss:', loss)
        return loss










